from playwright.sync_api import sync_playwright
import re
import time
import logging


def check_page(func):
    def wrap(self, *args, **kwargs):
        assert self.page
        return func(self, *args, **kwargs)

    return wrap


class Doubao:
    r"""## 使用方法：

    1. 先退出已有的豆包进程
    2. 然后cmd中执行：`C:\Users\Administrator\AppData\Local\Doubao\Application\Doubao.exe --remote-debugging-port=9222`
    """

    def __init__(self):
        self.connect_to_page()

    @check_page
    def ask(self, msg):
        print("问题：")
        print(msg)
        # 已发送消息数
        chat_yet_count = doubao.page.locator(
            "//div[contains(@data-testid, 'message_content') and contains(@class, 'justify-end')]"
        ).count()
        # 发送消息
        input_box = self.page.locator('textarea[placeholder*="发消息"]')
        parts = msg.split("\n")
        for part in parts:
            input_box.type(part)
            self.page.keyboard.press("Shift+Enter")
        input_box.type("\n")
        # 获得发送消息的定位
        while True:
            send_messages = doubao.page.locator(
                "//div[contains(@data-testid, 'message_content') and contains(@class, 'justify-end')]"
            )
            if send_messages.count() > chat_yet_count:
                send_message = send_messages.nth(-1)
                break
            time.sleep(0.1)
        #
        think_content = ""
        think_title = ""
        print("思考:")
        answer_locator = send_message.locator(
            "xpath=./following::div[@data-testid='message-block-container']"
        )
        think_title_locator = send_message.locator(
            "xpath=./following::span[contains(@class, 'think-block-title')]"
        )
        think_content_locator = send_message.locator(
            "xpath=./following::div[contains(@class, 'think-block-container')]"
        )
        while True:
            if not answer_locator.count():
                time.sleep(0.1)
                continue
            new_think_content = think_content_locator.nth(-1).inner_text()
            print(new_think_content.replace(think_content, ""), end="")
            # print("思考:", new_think_content)
            think_content = new_think_content
            time.sleep(0.1)
            try:
                think_title = think_title_locator.nth(-1).inner_text(timeout=30)
            except:
                pass
            if think_title == "已完成思考":
                print("")
                break
        answer = ""
        print("回答：")
        while True:
            answer_locator = doubao.page.locator(
                "//div[@data-testid='message_text_content']"
            ).nth(-1)
            new_answer = answer_locator.inner_text()
            print(new_answer.replace(answer, ""), end="")
            answer = new_answer
            if (
                answer_locator.locator(
                    "xpath=./following::div[contains(@class, 'suggest-message-list')]"
                ).count()
                > 0
            ):
                print("")
                break
            time.sleep(0.1)
        return {"question": msg, "think": think_content, "answer": answer}

    @check_page
    def new_chat(self):
        if self.page.url != "https://www.doubao.com/chat/":
            self.page.get_by_text("新对话").nth(0).click()
        assert self.page.url == "https://www.doubao.com/chat/"

    def connect_to_page(self):
        try:
            p = sync_playwright().start()
            browser = p.chromium.connect_over_cdp("http://localhost:9222")
            for page in browser.contexts[0].pages:
                if re.match(r"https://www.doubao.com/chat/", page.url):
                    self.page = page
                    break
        except:
            self.page = None


if __name__ == "__main__":
    doubao = Doubao()
    # doubao.connect_to_page()
    doubao.new_chat()
    # doubao.ask("https://blog.csdn.net/qq_26696715/article/details/123244639")
    result = doubao.ask("明天天气怎么样")
    doubao.ask("后天呢")
    pass
